from abc import ABC, abstractmethod
from pydantic import BaseModel

class BaseLLM(ABC):
    def __init__(self):
        pass

    @abstractmethod
    def get_max_input_tokens(self):
        pass

    @abstractmethod
    def get_max_output_tokens(self):
        pass

    @abstractmethod
    def get_context_token(self, llm_name: str, text: str):
        """
        args:
            llm_name: str, llm name
            text: str, text to be tokenized
        return:
            int, token count
        """
        try:
            # return len(self.tokenize_encode(text))
            return self.calculate_token(text)
        except Exception as e:
            return self.calculate_token(text)
        
    @abstractmethod
    def tokenize_encode(self, text: str):
        """
        args:
            text: str, text to be tokenized
        return:
            list, tokens
        """

    @abstractmethod
    def tokenize_decode(self, tokens: list):
        """
        args:
            tokens: list, tokens to be decoded
        return:
            str, text
        """

    def calculate_token(self, text: str):
        """
        args:
            text: str, text to be tokenized
        return:
            int, token count
        """
        per_chinese_token = 1
        per_english_token = 0.5
        per_other_token = 0.5
        total_tokens = 0
        for char in text:
            if char.isalpha():
                total_tokens += per_english_token
            elif '\u4e00' <= char <= '\u9fff':
                total_tokens += per_chinese_token
            else:
                total_tokens += per_other_token
                
        return total_tokens

    @abstractmethod
    def predict(self, text: str):
        """
        args:
            text: str, text to be predicted
        return:
            str, prediction result
        """

    @abstractmethod
    def apredict(self, text: str):
        """
        args:
            text: str, text to be predicted
        return:
            str, prediction result asynchronously
        """